# import typing

import openpyxl
from typing import Dict, List, Optional
import json


"""
    v1.0:
    根据测试集结构改变进行调整，
    新增了加载问题和类别1类别2

"""
class ExcelToDictLoader:
    def __init__(self, file_path: str):
        """
        初始化 Excel 到字典加载器
        :param file_path: Excel 文件路径
        """
        self.file_path = file_path
        # print(f"test3: {file_path}")
        self.workbook = None
        self.required_columns = ["问题","问题类别1","问题类别2","预期答案"]

    def load_workbook(self) -> bool:
        """加载 Excel 工作簿"""
        try:
            self.workbook = openpyxl.load_workbook(self.file_path, data_only=True)
            return True
        except Exception as e:
            print(f"[错误] 加载工作簿失败: {e}")
            return False

    def get_sheet_names(self) -> List[str]:
        """获取所有工作表名称"""
        if self.workbook is None and not self.load_workbook():
            return []
        return self.workbook.sheetnames

    def _locate_columns(self, sheet) -> Dict[str, int]:
        """定位目标列的索引位置"""
        column_map = {col: None for col in self.required_columns}
        # print(f"test1: {column_map}")
        # 扫描标题行(第一行)
        for cell in sheet[1]:
            cell_value = str(cell.value).strip() if cell.value else ""
            # print(f"test2: {cell_value}")
            for col in self.required_columns:
                if col in cell_value and column_map[col] is None:
                    column_map[col] = cell.column
                    break

        # 检查是否所有必要列都已找到
        missing = [col for col, idx in column_map.items() if idx is None]
        if missing:
            raise ValueError(f"缺少必要的列: {', '.join(missing)}")

        return column_map

    def load_to_dict(self, sheet_names: Optional[List[str]] = None) -> Dict[str, List[Dict[str, str]]]:
        """
        将 Excel 数据加载到字典结构
        :param sheet_names: 要加载的工作表名称列表(为空则加载所有工作表)
        :return: 字典结构 {工作表名: [记录字典, ...]}
        """
        if self.workbook is None and not self.load_workbook():
            return {}

        all_data = {}
        available_sheets = self.get_sheet_names()

        # 确定要处理的工作表
        sheets_to_load = sheet_names if sheet_names else available_sheets

        for sheet_name in sheets_to_load:
            if sheet_name not in available_sheets:
                print(f"[警告] 工作表 '{sheet_name}' 不存在，已跳过")
                continue

            sheet = self.workbook[sheet_name]
            try:
                col_map = self._locate_columns(sheet)
            except ValueError as e:
                print(f"[警告] 工作表 '{sheet_name}' 跳过: {e}")
                continue

            sheet_data = []
            for row in sheet.iter_rows(min_row=2):  # 跳过标题行

                if row[col_map["问题"] - 1].value != None and row[col_map["预期答案"] - 1].value != None and row[col_map["问题类别1"] - 1].value != "时间和天气": # 筛掉没有问题或答案的行，这里可以加问题类别的逻辑 TODO
                    # print(str(row[col_map["问题"] - 1].value))
                    try:
                        record = {
                            "问题": str(row[col_map["问题"] - 1].value).strip(),
                            "问题类别1": str(row[col_map["问题类别1"] - 1].value).strip(),
                            "问题类别2": str(row[col_map["问题类别2"] - 1].value).strip(),
                            "预期答案": str(row[col_map["预期答案"] - 1].value).strip(),
                            "行号": row[0].row  # 保留原始行号便于调试
                        }

                        # 跳过空问题或空答案的记录
                        if not record["问题"] or not record["预期答案"]:
                            continue

                        sheet_data.append(record)
                    except Exception as e:
                        print(f"[警告] 处理工作表 '{sheet_name}' 行 {row[0].row} 时出错: {e}")
                        continue

            all_data[sheet_name] = sheet_data

        return all_data

    def close(self):
        """关闭工作簿"""
        if self.workbook is not None:
            self.workbook.close()
            self.workbook = None

def load(sheet,filepath):
    """
    加载指定智能体（工作表）的数据
    :param sheet: 工作表
    :return: data_dict:指定工作表的所有数据，每条数据是一个字典，三个key：问题类型、问题、答案
    """
    # 1. 初始化加载器
    loader = ExcelToDictLoader(filepath)

    # 2. 查看可用工作表
    print("可用工作表:", ", ".join(loader.get_sheet_names()))

    # 3. 加载数据到字典 (可以指定特定工作表)
    # data_dict = loader.load_to_dict()  # 加载所有工作表
    data_dict = loader.load_to_dict([sheet,])  # 加载指定工作表(主要逻辑)

    # 4. 展示加载结果
    total_records = sum(len(records) for records in data_dict.values())
    print(f"\n成功加载 {total_records} 条有效记录，排除时间和天气类问题") # TODO：这里说明筛选条件

    for sheet_name, records in data_dict.items():
        print(f"\n工作表 '{sheet_name}' 有 {len(records)} 条有效记录:")
        for i, record in enumerate(records[:3], 1):  # 展示前3条作为示例
            print(f"  {i}. 问题类别1: {record['问题类别1']}")
            print(f"     问题类别2: {record['问题类别2']}")
            print(f"     问题: {record['问题']}")
            print(f"     答案: {record['预期答案']}")
        print("     ...")

    # 5. 关闭加载器
    loader.close()


    # # 6. 使用字典数据进行测试
    # # 示例测试函数
    # def run_tests(test_data: Dict[str, List[Dict[str, str]]]):
    #     for sheet_name, records in test_data.items():
    #         print(f"\n正在测试工作表: {sheet_name}")
    #         for record in records:
    #             # 这里替换为实际的测试逻辑
    #             test_result = "通过" if len(record["问题"]) < 50 else "失败"
    #             print(f"[{test_result}] 行{record['行号']}: {record['问题'][:30]}...")
    #
    #
    # # 运行测试
    # run_tests(data_dict)

    return data_dict

def load_question_only(sheet="聊天助手"):
    """
    只加载问题
    :param sheet:
    :return:
    """
    print("loading question...")
    data_sheet=load(sheet)
    question=[]
    for i in data_sheet[sheet]:
        question.append(i["问题"])
    print("loading question complete")
    return question

def load_reference_only(sheet="聊天助手"):
    """
        只加载参考答案
        :param sheet:
        :return:
        """
    print("loading reference...")
    data_sheet = load(sheet,filepath="../data/同方智能体测试集 250507.xlsx")
    reference = []
    for i in data_sheet[sheet]:
        reference.append(i["预期答案"])
    print("loading reference complete")
    return reference
def load_answer(file_answer="聊天助手_answers_20250430_094618.json"):
    """
    从文件加载智能体的回答
    :param file:回答文件
    :return:
    """
    # 读取 JSON 文件
    with open(f'../data/{file_answer}', 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 提取 answer 字段组成数组
    answers = [item['answer'] for item in data]
    return answers


def load_question_and_categories(sheet="聊天助手", filepath="../data/同方智能体测试集 250507.xlsx"):
    """
    加载问题及其对应的两个类别信息
    :param sheet: 工作表名称
    :param filepath: 文件路径
    :return: 包含问题、类别1和类别2的字典列表
    """
    print("loading questions with categories...")
    data_sheet = load(sheet, filepath=filepath)
    questions_with_categories = []

    for item in data_sheet[sheet]:
        questions_with_categories.append({
            "question": item["问题"],
            "category1": item["问题类别1"],  # 使用get方法防止键不存在
            "category2": item.get("问题类别2", "")
        })

    print("loading questions with categories complete")
    return questions_with_categories


def load_categories(sheet="聊天助手", filepath="../data/同方智能体测试集 250507.xlsx"):
    print("loading questions with categories...")
    data_sheet = load(sheet, filepath)

    # 调试：打印第一条记录的所有键
    if data_sheet[sheet]:
        print("可用字段:", data_sheet[sheet][0].keys())

    questions_with_categories = []
    # ...其余代码不变

if __name__=="__main__":
    data=load_question_and_categories()
    # print(data)
    # load_categories()
    for i,row in enumerate(data):
        print(i,row)
    # help(typing)
